package main

import (
	"fmt"
	"io"
	"log"
	"net/http"
	"net/url"
	"os"
	"path/filepath"
	"strconv"
	"strings"
	"time"

	"github.com/gin-gonic/gin"
)

func main() {
	// 设置为生产模式
	if os.Getenv("GIN_MODE") != "debug" {
		gin.SetMode(gin.ReleaseMode)
	}

	// 创建默认路由引擎
	r := gin.Default()

	// 代理下载接口 - 使用GET方法
	r.GET("/download/proxy", proxyHandler)

	// 启动服务器
	port := "3333"
	log.Printf("服务器启动在端口 %s\n", port)
	log.Fatal(r.Run(":" + port))
}

// 获取文件大小（通过HEAD请求）
func getFileSize(url string) (int64, error) {
	req, err := http.NewRequest("HEAD", url, nil)
	if err != nil {
		return 0, err
	}

	// 设置请求头（部分服务器需要User-Agent才能响应HEAD请求）
	req.Header.Set("User-Agent", "Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/91.0.4472.124 Safari/537.36")

	client := &http.Client{
		Timeout: 10 * time.Second, // HEAD请求超时不宜过长
	}

	resp, err := client.Do(req)
	if err != nil {
		return 0, err
	}
	defer resp.Body.Close()

	// 若服务器不支持HEAD，尝试用GET请求的响应头（只读取头，不下载内容）
	if resp.StatusCode == http.StatusMethodNotAllowed {
		req, err = http.NewRequest("GET", url, nil)
		if err != nil {
			return 0, err
		}
		req.Header.Set("Range", "bytes=0-0") // 只请求第一个字节，获取文件大小
		resp, err = client.Do(req)
		if err != nil {
			return 0, err
		}
		defer resp.Body.Close()
	}

	if resp.StatusCode < 200 || resp.StatusCode >= 300 {
		return 0, fmt.Errorf("远程服务器返回 %d 状态码", resp.StatusCode)
	}

	contentLength := resp.Header.Get("Content-Length")
	if contentLength == "" {
		return 0, fmt.Errorf("未获取到Content-Length")
	}

	return strconv.ParseInt(contentLength, 10, 64)
}

const (
	smallFileThreshold = 50 * 1024 * 1024   // 小文件阈值：50MB
	smallClientTimeout = 60 * time.Second   // 小文件客户端超时
	largeClientTimeout = 3600 * time.Second // 大文件客户端超时
	smallWriteTimeout  = 120 * time.Second  // 小文件写入超时
	largeWriteTimeout  = 3600 * time.Second // 大文件写入超时
	defaultTimeout     = 30 * time.Second   // 默认超时（无法获取文件大小时）
)

// 代理下载处理函数
func proxyHandler(c *gin.Context) {
	// 从URL参数获取远程资源URL
	targetURL := c.Query("url")

	if targetURL == "" {
		c.Status(http.StatusBadRequest)
		c.Writer.WriteString("错误: 缺少URL参数")
		return
	}

	// 验证URL有效性
	parsedURL, err := url.ParseRequestURI(targetURL)
	if err != nil {
		c.Status(http.StatusBadRequest)
		c.Writer.WriteString(fmt.Sprintf("错误: 无效的URL - %v", err))
		return
	}

	fileSize, err := getFileSize(parsedURL.String())
	if err != nil {
		log.Printf("获取文件大小失败，使用默认超时: %v", err)
		fileSize = -1 // 标记为未知大小
	}

	// 步骤2：根据文件大小动态设置超时
	clientTimeout, _ := getDynamicTimeouts(fileSize)

	// 创建HTTP客户端并设置超时
	client := &http.Client{Timeout: clientTimeout}

	// 创建请求
	req, err := http.NewRequest("GET", targetURL, nil)
	if err != nil {
		c.Status(http.StatusInternalServerError)
		c.Writer.WriteString(fmt.Sprintf("错误: 创建请求失败 - %v", err))
		return
	}

	// 设置常见请求头，模拟浏览器行为
	req.Header.Set("User-Agent", "Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/91.0.4472.124 Safari/537.36")
	req.Header.Set("Accept", "text/html,application/xhtml+xml,application/xml;q=0.9,image/webp,*/*;q=0.8")
	req.Header.Set("Accept-Language", "zh-CN,zh;q=0.9,en;q=0.8")

	// 发送请求
	resp, err := client.Do(req)
	if err != nil {
		c.Status(http.StatusInternalServerError)
		c.Writer.WriteString(fmt.Sprintf("错误: 请求远程资源失败 - %v", err))
		return
	}
	defer resp.Body.Close()

	// 检查响应状态
	if resp.StatusCode != http.StatusOK {
		c.Status(http.StatusInternalServerError)
		c.Writer.WriteString(fmt.Sprintf("错误: 远程服务器返回 %d 状态码", resp.StatusCode))
		return
	}

	// 尝试从Content-Disposition头或URL中提取文件名
	filename := extractFilename(targetURL, resp.Header.Get("Content-Disposition"))

	// 设置响应头
	var contentLength string
	if fileSize > 0 {
		contentLength = strconv.FormatInt(fileSize, 10)
	} else {
		contentLength = resp.Header.Get("Content-Length")
	}

	if contentLength != "" {
		// 确保设置Content-Length，不使用分块传输
		c.Writer.Header().Set("Content-Length", contentLength)
		// 移除可能存在的分块传输头
		c.Writer.Header().Del("Transfer-Encoding")
	} else {
		// 确实无法获取大小，才使用分块传输
		c.Writer.Header().Set("Transfer-Encoding", "chunked")
	}
	c.Writer.Header().Set("Content-Disposition", fmt.Sprintf("attachment; filename=\"%s\"", filename))
	if contentType := resp.Header.Get("Content-Type"); contentType != "" {
		if contentType == "image/webp" {
			contentType = "image/jpg"
		}
		c.Writer.Header().Set("Content-Type", contentType)
	}

	// 流式传输响应内容到客户端
	_, err = io.Copy(c.Writer, resp.Body)
	if err != nil {
		log.Printf("错误: 复制响应内容失败 - %v", err)
	}
}

// 从URL和Content-Disposition头提取文件名
func extractFilename(urlStr, contentDisposition string) string {
	// 优先从Content-Disposition中提取
	if contentDisposition != "" {
		// 格式可能是 "attachment; filename=example.zip"
		if strings.Contains(contentDisposition, "filename=") {
			parts := strings.Split(contentDisposition, "filename=")
			if len(parts) > 1 {
				filename := strings.Trim(parts[1], `" '`)
				if filename != "" {
					return filename
				}
			}
		}
	}

	// 从URL路径中提取
	u, err := url.Parse(urlStr)
	if err != nil {
		return "file.bin"
	}

	// 获取路径部分并提取最后一段
	path := u.Path
	if path == "" {
		return "file.bin"
	}

	// 提取文件名部分
	filename := filepath.Base(path)
	if filename == "" || filename == "." || filename == ".." {
		return "file.bin"
	}

	return filename
}

func getDynamicTimeouts(fileSize int64) (clientTimeout, writeTimeout time.Duration) {
	if fileSize == -1 {
		// 未知大小，使用默认超时
		return defaultTimeout, defaultTimeout
	}

	if fileSize < smallFileThreshold {
		// 小文件
		return smallClientTimeout, smallWriteTimeout
	} else {
		// 大文件
		return largeClientTimeout, largeWriteTimeout
	}
}
